#!/usr/bin/env python

from argparse import ArgumentParser
import struct

def parse_header(data: bytes):
    print("Data size: ", len(data))
    magic = struct.unpack("<I", data[:4])[0]
    header_len = struct.unpack("<I", data[4:8])[0]
    reserved = struct.unpack("<I", data[20:24])[0]
    image_type = struct.unpack("<H", data[24:26])[0]
    storage_type = data[26]
    signature_type = data[27]
    load_address = struct.unpack("<I", data[28:32])[0]
    total_file_size = struct.unpack("<I", data[32:36])[0]
    maximum_file_size = struct.unpack("<I", data[36:40])[0]
    content_offset = struct.unpack("<I", data[40:44])[0]
    signature_length = struct.unpack("<I", data[44:48])[0]
    jump_offset = struct.unpack("<I", data[48:52])[0]
    flags = struct.unpack("<I", data[52:56])[0]
    print("==================")
    print("Magic: 0x%08x" % magic)
    print("Header Length: %d" % header_len)
    print("Reserved: %d" % reserved)
    if image_type == 0:
        print("Image Type: NONE")
    elif image_type == 1:
        print("Image Type: ARM Bootloader")
    elif image_type == 2:
        print("Image Type: ARM External Bootloader")
    elif image_type == 10:
        print("Image Type: Root Certificate")
    elif image_type == 256:
        print("Image Type: Primary MAUI")
    elif image_type == 264:
        print("Image Type: VIVA")
    elif image_type == 769:
        print("Image Type: SECURE_RO_ME")
    else:
        print("Image Type: Unknown(%d)" % image_type)

    if storage_type == 0:
        print("Storage Type: NONE")
    elif storage_type == 1:
        print("Storage Type: NOR Flash")
    elif storage_type == 2:
        print("Storage Type: NAND Sequential Flash")
    elif storage_type == 3:
        print("Storage Type: NAND TTBL")
    elif storage_type == 4:
        print("Storage Type: NAND FDM50")
    elif storage_type == 5:
        print("Storage Type: EMMC Boot Region")
    elif storage_type == 6:
        print("Storage Type: EMMC Data Region")
    elif storage_type == 7:
        print("Storage Type: Serial Flash")
    else:
        print("Storage Type: Unknown(%d)" % storage_type)

    if signature_type == 0:
        print("Signature Type: NONE")
    elif signature_type == 1:
        print("Signature Type: PHASH")
    elif signature_type == 2:
        print("Signature Type: SINGLE")
    elif signature_type == 3:
        print("Signature Type: SINGLE + PHASH")
    elif signature_type == 4:
        print("Signature Type: MULTI")
    else:
        print("Signature Type: Unknown(%d)" % signature_type)

    print("Load Address: 0x%08x" % load_address)
    print("Total File Size: %d (0x%x)" % (total_file_size, total_file_size))
    print("Maximum File Size: %d (0x%x)" % (maximum_file_size, maximum_file_size))
    print("Content Offset: %d (0x%x)" % (content_offset, content_offset))
    print("Signature Length: %d (0x%x)" % (signature_length, signature_length))
    print("Jump Offset: 0x%08x" % jump_offset)
    print("Flags 0x%x" % flags)
    print("")

    return dict({
        "image_type": image_type,
        "storage_type": storage_type,
        "signature_type": signature_type,
        "load_address": load_address,
        "total_file_size": total_file_size,
        "content_offset": content_offset,
        "signature_len": signature_length,
        "jump_offset": jump_offset,
        "flags": flags
    })

def parse_headers(data:bytes):
    header_offset = data.find(b'\x4d\x4d\x4d\x01')
    print("header offset: ", header_offset)
    parse_header(data[header_offset:])

def main():
    parser = ArgumentParser()
    parser.add_argument("file")
    args = parser.parse_args()

    with open(args.file, "rb") as f:
        data = f.read()
        parse_headers(data)
        f.close()


if __name__ == "__main__":
    main()